#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <linux/types.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>
#include <errno.h>


typedef enum {
    OPENSM_NULL = 0,
    OPENSM_GENKEY,
    OPENSM_ENCRYPT,
    OPENSM_DECRYPT,
} OPENSM_FUNC;



typedef enum {
    SM4_CRYPTO_ENCRYPT = 0,
    SM4_CRYPTO_DECRYPT,
} SM4_CRYPTO_E;

typedef enum
{
    SM4_MODE_ECB = 0,   // ECB Mode
    SM4_MODE_CBC,   // CBC Mode
    SM4_MODE_CFB,   // CFB Mode
    SM4_MODE_OFB,   // OFB Mode
    SM4_MODE_CTR    // CTR Mode
} SM4_MODE_E;


typedef struct _st_sm2_key
{
    uint8_t priKey[32];
    uint8_t pubKey[65];
}SM2_KEY;

typedef struct _st_sm2_encrypt
{
    uint8_t M[256];
    uint32_t MByteLen;
    uint8_t pubKey[65];
    uint8_t order;
    uint8_t C[512];
    uint32_t CByteLen;
}SM2_ENCRYPT;

typedef struct _st_sm2_decrypt
{
    uint8_t C[512];
    uint32_t CByteLen;
    uint8_t priKey[32];
    uint8_t order;
    uint8_t M[256];
    uint32_t MByteLen;
}SM2_DECRYPT;

typedef struct _st_hash_context
{   
    uint8_t first_update_flag;
    uint8_t finish_flag;
    uint8_t hash_buffer[64];
    uint32_t total[2];
}SM3_HASH_CONTEXT;
typedef struct _st_hash
{
    uint8_t message[128];
    uint32_t byteLen;
    uint8_t digest[32];
}SM3_HASH;


typedef struct _st_crypto_init
{
    SM4_MODE_E mode;
    SM4_CRYPTO_E crypto;
    uint8_t key[128];
    uint8_t iv[128];
}SM4_INIT;

typedef struct _st_crypto_
{
    uint8_t in[128];
    uint8_t out[128];
    uint32_t byteLen;
}SM4_CRYPTO;



#define SCTO_SM2_KEYGET            _IOR('k', 0, SM2_KEY *)
#define SCTO_SM2_ENCRYPT           _IOR('k', 1, SM2_ENCRYPT *)
#define SCTO_SM2_DECRYPT           _IOR('k', 2, SM2_DECRYPT *)
#define SCTO_SM3_INIT              _IOR('k', 9,  SM3_HASH_CONTEXT *)
#define SCTO_SM3_HASH              _IOR('k', 12, SM3_HASH *)
#define SCTO_SM4_INIT              _IOR('k', 20,  SM4_INIT *)
#define SCTO_SM4_CRYPTO            _IOR('k', 21,  SM4_CRYPTO *)

int enc_file(char *input_file, char *output_file, char *prikey_file, char *pubkey_file);
int dec_file(char *input_file, char *output_file, char *prikey_file);
int gen_key_pair(char *prikey_file, char *pubkey_file);

#define VERBOSE0(fmt, ...) do {		 \
		printf(fmt, ## __VA_ARGS__);	\
	} while (0)

static void usage ()
{
	VERBOSE0 ("OPENSM Tool.\n"
	"Usage: \n"
    "-h, --help         Print usage information\n"
    "-i, --infile		File input\n"
    "-o, --outfile		File output\n"
    "-e, --encrypt		Encryption\n"
    "-d, --decrypt		Dncryption\n"
    "-g, --genkey		Generate key pair\n"
    "-p, --pubkey		Public key file\n"
    "-k, --prikey		Private key file\n"
    "EXAMPLE:\n"
    "1. Generating your SM2 key pair:\n"
    "opensm -g -k $(path)/your.prikey -p $(path)/your.pubkey\n"
    "2. Encrypting your file for local storage:\n"
    "opensm -e -p $(path)/your.pubkey -i $(path)/plaintext -o $(path)/ciphertext \n"
    "3. Encrypting your file (unencrypted) for network transmission:\n"
    "opensm -e -p $(path)/his.pubkey -i $(path)/plaintext -o $(path)/ciphertext \n"
    "4. Encrypting your file (encrypted) for network transmission:\n"
    "opensm -e -k $(path)/your.prikey -p $(path)/his.pubkey -i $(path)/ciphertext -o $(path)/new_ciphertext \n"
    "5. Decrypting the encrypted file:\n"
    "opensm -d -k $(path)/your.prikey -i $(path)/ciphertext -o $(path)/plaintext \n"
    );
}


int main (int argc, char* argv[])
{
    char *input_file = NULL;
    char *output_file = NULL;
	char *prikey_file = NULL;
	char *pubkey_file = NULL;

    OPENSM_FUNC my_func = OPENSM_NULL;

    int cmd;

	while (1) {
		int option_index = 0;
		static struct option long_options[] = {
			{ "encrypt",	 	no_argument, 	        NULL, 	'e' },
            { "decrypt",	 	no_argument, 	        NULL, 	'd' },
            { "genkey",	 	    no_argument, 	        NULL, 	'g' },
			{ "help",	 	    no_argument, 		    NULL, 	'h' },
			{ "infile",	    	required_argument, 	    NULL, 	'i' },
            { "outfile",		required_argument, 	    NULL, 	'o' },
			{ "prikey",		    required_argument, 	    NULL, 	'k' },
			{ "pubkey",		    required_argument, 	    NULL, 	'p' },
		};
		cmd = getopt_long (argc, argv, "i:o:w:k:p:hedg",
						   long_options, &option_index);

		if (cmd == -1) { /* all params processed ? */
			break;
		}

		switch (cmd) {
		case 'h':   /* help */
			usage ();
			exit (EXIT_SUCCESS);
		case 'i':  
			input_file = optarg;
			break;
		case 'o':
			output_file = optarg; 
			break;
		case 'e':
            my_func = OPENSM_ENCRYPT;
            break;
        case 'd':
            my_func = OPENSM_DECRYPT;
            break;
        case 'g':
            my_func = OPENSM_GENKEY;
            break;
		case 'p':	  
			pubkey_file = optarg;
			break;
		case 'k':	  
			prikey_file = optarg;
			break;
		default:
			usage ();
			exit (EXIT_FAILURE);
		}
	}

    int ret = 0;

    //printf("pubkey_file :%s\n", pubkey_file);
    //printf("prikey_file :%s\n", prikey_file);
    //printf("input_file :%s\n", input_file);
    //printf("output_file :%s\n", output_file);

    if(my_func == OPENSM_GENKEY){
        if((prikey_file== NULL) || (pubkey_file== NULL)){
            usage ();
			exit (EXIT_FAILURE);
        } else {
            ret = gen_key_pair(prikey_file, pubkey_file);
        }
        
    } else if (my_func == OPENSM_ENCRYPT){
        if((pubkey_file== NULL) || (input_file == NULL) || (output_file == NULL)){
            usage ();
			exit (EXIT_FAILURE);
        } else {
            ret = enc_file(input_file, output_file, prikey_file, pubkey_file);
        }
    } else if (my_func == OPENSM_DECRYPT){
        if((prikey_file == NULL) || (input_file == NULL) || (output_file == NULL)){
            usage ();
			exit (EXIT_FAILURE);
        } else {
            ret = dec_file(input_file, output_file, prikey_file);
        }
    } else {
        usage ();
		exit (EXIT_SUCCESS);
    }

    if(ret != 0){
        printf("opensm runs failed!\n");
        usage ();
		exit (EXIT_FAILURE);
    }
    else{
        printf("opensm runs successfully!\n");
    }

    return 0;
}



int gen_key_pair(char *prikey_file, char *pubkey_file){
    int fd = open("/dev/dma", O_RDWR);
    if( fd < 0 ){
        printf("open /dev/dma failed\n");
        return -1;
    }

    int ret = 0;

    FILE *fp;

    SM2_KEY val;
    ret = ioctl(fd, SCTO_SM2_KEYGET, &val);
    if( ret != 0 ){
        printf("SCTO_SM2_KEYGET failed with ret:%d\n",ret);
        close(fd);
        return -1; 
    }

    if( (fp=fopen(pubkey_file, "wb")) == NULL ){  //以二进制方式打开
        puts("Fail to open pubkey_file!");
        return -1;
    }
    fwrite(val.pubKey, 1, sizeof(val.pubKey), fp);
    fclose(fp);

    if( (fp=fopen(prikey_file, "wb")) == NULL ){  //以二进制方式打开
        puts("Fail to open prikey_file!");
        return -1;
    }
    fwrite(val.priKey, 1, sizeof(val.priKey), fp);
    fclose(fp);

    close(fd);
    return 0;
}


int enc_file(char *input_file, char *output_file, char *prikey_file, char *pubkey_file){
    int fd = open("/dev/dma", O_RDWR);
    if( fd < 0 ){
        printf("open /dev/dma failed\n");
        return -1;
    }

    int ret = 0;

    FILE *fp;

    SM2_ENCRYPT     sm2_encrypt;
    SM2_DECRYPT     sm2_decrypt;
    SM2_KEY         val;
    SM2_KEY         val2;
    SM3_HASH        sm3_file;
    SM4_INIT        sm4_init;
    SM4_CRYPTO      sm4_enc_crypto;

    if( (fp=fopen(pubkey_file, "rb")) == NULL ){  //以二进制方式打开
        puts("Fail to open pubkey_file!");
        return -1;
    }
    fseek(fp, 0, SEEK_END);        
    uint32_t pubkey_size = ftell(fp); 
    if(pubkey_size != 65){
        puts("Wrong pubkey_file!");
        fclose(fp);
        return -1;
    } else {
        fseek(fp, 0, SEEK_SET);
        fread(val.pubKey, 1, sizeof(val.pubKey), fp);
        fclose(fp); 
    }
    

    if( (fp=fopen(input_file, "rb")) == NULL ){  //以二进制方式打开
        puts("Fail to open input_file!");
        return -1;
    }
    fseek(fp, 0, SEEK_END);        //指针：移动到文件尾部
    uint32_t f_size = ftell(fp);   //返回指针偏离文件头的位置（即文件中字符个数）
    uint32_t block_num, rendent_num, padding_num;
    uint8_t *temp_data;


    if(prikey_file != NULL){
        if(f_size <= 129){
            puts("The input_file is wrong!");
            return -1;
        }
        f_size = f_size;
        
        temp_data = (uint8_t *)malloc(sizeof(uint8_t)*(f_size));
        memset(temp_data, 0, f_size);
        fseek(fp, 0, SEEK_SET);  //指针：移动到文件头部
        fread(temp_data, 1, f_size, fp);
        fclose(fp);


        if( (fp=fopen(prikey_file, "rb")) == NULL ){  //以二进制方式打开
            puts("Fail to open prikey_file!");
            return -1;
        }
        fseek(fp, 0, SEEK_END);        
        uint32_t prikey_size = ftell(fp); 
        if(prikey_size != 32){
            puts("Wrong prikey_file!");
            fclose(fp);
            return -1;
        } else {
            fseek(fp, 0, SEEK_SET);
            fread(val.priKey, 1, sizeof(val.priKey), fp);
            fclose(fp);
        }

        //解密密钥
        memset(&sm2_decrypt, 0, sizeof(sm2_decrypt));
        memcpy(sm2_decrypt.C, temp_data, 129);
        sm2_decrypt.CByteLen = 129;
        memcpy(sm2_decrypt.priKey, val.priKey, sizeof(val.priKey));
        sm2_decrypt.order = 0;
        ret = ioctl(fd, SCTO_SM2_DECRYPT, &sm2_decrypt); 
        if( ret != 0 ){
            printf("SCTO_SM2_DECRYPT failed with ret:%d\n",ret);
            close(fd);
            return -1; 
        }

        //加密文件密钥
        memset(&sm2_encrypt, 0, sizeof(sm2_encrypt));
        memcpy(sm2_encrypt.M, sm2_decrypt.M, 32);
        sm2_encrypt.MByteLen = 32;
        memcpy(sm2_encrypt.pubKey, val.pubKey, sizeof(val.pubKey));
        sm2_encrypt.order = 0;

        ret = ioctl(fd, SCTO_SM2_ENCRYPT, &sm2_encrypt);
        if( ret != 0 ){
            printf("SCTO_SM2_ENCRYPT failed with ret:%d\n",ret);
            close(fd);
            return -1; 
        }

        memcpy(temp_data, sm2_encrypt.C, 129);


        if( (fp=fopen(output_file, "wb")) == NULL ){  //以二进制方式打开
            puts("Fail to open output_file!");
            return -1;
        }
        fwrite(temp_data, 1, f_size, fp);

        fclose(fp);

    } else {
        block_num = f_size / 128;
        rendent_num = f_size - block_num * 128;
        padding_num = 0;
        if(rendent_num != 0){
            padding_num = 128 - rendent_num;
            block_num = block_num + 1;
        }
        temp_data = (uint8_t *)malloc(sizeof(uint8_t)*(f_size+padding_num));
        memset(temp_data, 0, f_size+padding_num);
        fseek(fp, 0, SEEK_SET);  //指针：移动到文件头部
        fread(temp_data, 1, f_size, fp);
        fclose(fp);

        ret = ioctl(fd, SCTO_SM2_KEYGET, &val2);
        if( ret != 0 ){
            printf("SCTO_SM2_KEYGET failed with ret:%d\n",ret);
            close(fd);
            return -1; 
        }
        //sm3
        bzero(&sm3_file, sizeof(sm3_file));
        memcpy(sm3_file.message, val2.priKey, 32);
        memcpy(sm3_file.message+32, temp_data, 96);
        sm3_file.byteLen = 128;

        ret = ioctl(fd, SCTO_SM3_HASH, &sm3_file);
        if( ret != 0){
            close(fd);
            return -1;
        }

        //加密文件密钥
        memset(&sm2_encrypt, 0, sizeof(sm2_encrypt));
        memcpy(sm2_encrypt.M, sm3_file.digest, 32);
        sm2_encrypt.MByteLen = 32;
        memcpy(sm2_encrypt.pubKey, val.pubKey, sizeof(val.pubKey));
        sm2_encrypt.order = 0;

        ret = ioctl(fd, SCTO_SM2_ENCRYPT, &sm2_encrypt);
        if( ret != 0 ){
            printf("SCTO_SM2_ENCRYPT failed with ret:%d\n",ret);
            close(fd);
            return -1; 
        }

        //初始化sm4
        bzero(&sm4_init, sizeof(sm4_init));
        sm4_init.mode = SM4_MODE_OFB;
        sm4_init.crypto = SM4_CRYPTO_ENCRYPT;
        memcpy(sm4_init.key, sm3_file.digest, 16);
        memcpy(sm4_init.iv, sm3_file.digest +16, 16);
        ret = ioctl(fd, SCTO_SM4_INIT, &sm4_init);
        if( ret != 0){
        	printf("SCTO_SM4_INIT failed with ret:%d\n",ret);
        	close(fd);
        	return -1;
        }

        //保存加密后的sm4密钥和密文
        if( (fp=fopen(output_file, "wb")) == NULL ){  //以二进制方式打开
            puts("Fail to open ciphertext file1!");
            return -1;
        }
        fwrite(sm2_encrypt.C, 1, sm2_encrypt.CByteLen, fp);


        for (int i = 0; i < block_num; i++){
            bzero(&sm4_enc_crypto, sizeof(sm4_enc_crypto));
            if(prikey_file != NULL){
                memcpy(sm4_enc_crypto.in, temp_data + 129 + 128 * i, 128);
            } else {
                memcpy(sm4_enc_crypto.in, temp_data + 128 * i, 128);
            }
            sm4_enc_crypto.byteLen = 128;
            ret = ioctl(fd, SCTO_SM4_CRYPTO, &sm4_enc_crypto);
            if( ret != 0){
            	printf("SCTO_SM4_CRYPTO2 failed with ret:%d\n",ret);
            	close(fd);
            	return -1;
            }
            if(i == block_num -1){
                fwrite(sm4_enc_crypto.out, 1, rendent_num, fp);
            } else {
                fwrite(sm4_enc_crypto.out, 1, 128, fp);
            }

        }

        fclose(fp);
    }
    
    

    close(fd);

    return 0;
}



int dec_file(char *input_file, char *output_file, char *prikey_file){
    int fd = open("/dev/dma", O_RDWR);
    if( fd < 0 ){
        printf("open /dev/dma failed\n");
        return -1;
    }

    int ret = 0;


    //读取密文
    FILE *fp;

    SM2_KEY val;
    SM2_DECRYPT sm2_decrypt;
    SM4_INIT sm4_init;
    SM4_CRYPTO sm4_dec_crypto;

    if( (fp=fopen(input_file, "rb")) == NULL ){  //以二进制方式打开
        puts("Fail to open input_file!");
        exit(0);
    }

    fseek(fp, 0, SEEK_END);        //指针：移动到文件尾部
    uint32_t f_size = ftell(fp) - 129;   //返回指针偏离文件头的位置（即文件中字符个数）
    uint32_t block_num = f_size / 128;
    uint32_t rendent_num = f_size - block_num * 128;
    uint32_t padding_num = 0;
    if(rendent_num != 0){
        padding_num = 128 - rendent_num;
        block_num = block_num + 1;
    }

    uint8_t *temp_data = (uint8_t *)malloc(sizeof(uint8_t)*(f_size+padding_num + 129));
    memset(temp_data, 0, f_size+padding_num + 129);
    fseek(fp, 0, SEEK_SET);  //指针：移动到文件头部
    fread(temp_data, 1, f_size+129, fp);
    fclose(fp);

    //读取私钥
    if( (fp=fopen(prikey_file, "rb")) == NULL ){  //以二进制方式打开
        puts("Fail to open prikey_file!");
        return -1;
    }
    fseek(fp, 0, SEEK_END);        
    uint32_t prikey_size = ftell(fp); 
    if(prikey_size != 32){
        puts("Wrong prikey_file!");
        fclose(fp);
        return -1;
    } else {
        fseek(fp, 0, SEEK_SET);
        fread(val.priKey, 1, sizeof(val.priKey), fp);
        fclose(fp);
    }

    //解密密钥
    
    memset(&sm2_decrypt, 0, sizeof(sm2_decrypt));
    memcpy(sm2_decrypt.C, temp_data, 129);
    sm2_decrypt.CByteLen = 129;
    memcpy(sm2_decrypt.priKey, val.priKey, sizeof(val.priKey));
    sm2_decrypt.order = 0;
    ret = ioctl(fd, SCTO_SM2_DECRYPT, &sm2_decrypt); 
    if( ret != 0 ){
        printf("SCTO_SM2_DECRYPT failed with ret:%d\n",ret);
        close(fd);
        return -1; 
    }

    //解密文件
    
    bzero(&sm4_init, sizeof(sm4_init));
    sm4_init.mode = SM4_MODE_OFB;
    sm4_init.crypto = SM4_CRYPTO_DECRYPT;
    memcpy(sm4_init.key, sm2_decrypt.M, 16);
    memcpy(sm4_init.iv, sm2_decrypt.M +16, 16);
    ret = ioctl(fd, SCTO_SM4_INIT, &sm4_init);
    if( ret != 0){
    	printf("SCTO_SM4_INIT2 failed with ret:%d\n",ret);
    	close(fd);
    	return -1;
    }

    if( (fp=fopen(output_file, "wb")) == NULL ){  //以二进制方式打开
        puts("Fail to open output_file!");
        exit(0);
    }
    
    
    for (int i = 0; i < block_num; i++){
        bzero(&sm4_dec_crypto, sizeof(sm4_dec_crypto));
        memcpy(sm4_dec_crypto.in, temp_data +129 + 128*i, 128);
        //memcpy(sm4_dec_crypto.in, std_in, strlen(std_in));
        sm4_dec_crypto.byteLen = 128;
        ret = ioctl(fd, SCTO_SM4_CRYPTO, &sm4_dec_crypto);
        if( ret != 0){
        	printf("SCTO_SM4_CRYPTO2 failed with ret:%d\n",ret);
        	close(fd);
        	return -1;
        }
        if(i == block_num -1){
            fwrite(sm4_dec_crypto.out, 1, rendent_num, fp);
        } else {
            fwrite(sm4_dec_crypto.out, 1, 128, fp);
        }
    }

    fclose(fp);

    close(fd);

    return 0;
}